/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Description: Stub cpp generator.
 */
#include "datasystem/common/rpc/plugin_generator/zmq_rpc_generator.h"

namespace datasystem {
void ZmqRpcGenerator::CreateStubCpp(const google::protobuf::FileDescriptor &file,
                                    compiler::GeneratorContext *generatorCtx) const
{
    std::unique_ptr<io::ZeroCopyOutputStream> outputFile(generatorCtx->Open(fileName + ".stub.rpc.pb.cc"));
    io::Printer printer(outputFile.get(), '$');
    std::map<std::string, std::string> vars;
    vars["fileName"] = fileName;
    vars["fileFullName"] = file.name();
    std::string impl =
        "// Generated by the protocol buffer compiler.  DO NOT EDIT!\n"
        "// source: $fileFullName$\n"
        "#include \"$fileName$.service.rpc.pb.h\"\n"
        "#include \"$fileName$.stub.rpc.pb.h\"\n"
        "#include \"datasystem/common/rpc/rpc_unary_client_impl.h\"\n"
        "#include \"datasystem/common/rdma/urma_manager_wrapper.h\"\n"
        "#include \"datasystem/common/rpc/zmq/zmq_client_stream_base.h\"\n"
        "#include \"datasystem/common/rpc/zmq/zmq_stub_impl.h\"\n";
    printer.Print(vars, impl.c_str());
    InstantiateTemplate(printer, file);
    printer.PrintRaw(namespaceBegin);
    for (auto i = 0; i < file.service_count(); ++i) {
        auto *svc = file.service(i);
        const std::string &svcName = svc->name();
        const std::string stub = svcName + "_ZmqStub";
        ImplementGenericStubApiDef(printer, *svc, PREFIX);
        GenerateInitMethodMapDef(printer, *svc, PREFIX, stub);
        ImplementStubApiDef(printer, *svc, PREFIX);
    }
    printer.PrintRaw(namespaceEnd);
}

void ZmqRpcGenerator::InstantiateService(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc) const
{
    std::set<std::pair<std::string, std::string>> streamClass;
    std::set<std::string> writerClass;
    std::set<std::string> readerClass;
    std::map<std::string, std::string> vars;
    for (int j = 0; j < svc.method_count(); j++) {
        auto *method = svc.method(j);
        const std::string &inputName = method->input_type()->name();
        const std::string &outputName = method->output_type()->name();
        vars["inputTypeName"] = inputName;
        vars["outputTypeName"] = outputName;
        if (method->client_streaming() && method->server_streaming()) {
            auto ele = std::make_pair(inputName, outputName);
            if (streamClass.count(ele) == 0) {
                printer.Print(vars, "template class ClientWriterReader<$inputTypeName$, $outputTypeName$>;\n");
                streamClass.insert(ele);
            }
        } else if (method->client_streaming()) {
            if (writerClass.count(inputName) == 0) {
                printer.Print(vars, "template class ClientWriter<$inputTypeName$>;\n");
                printer.Print(vars, "template Status ClientWriter<$inputTypeName$>::Read($outputTypeName$ &);\n");
                writerClass.insert(inputName);
            }
        } else if (method->server_streaming()) {
            if (readerClass.count(outputName) == 0) {
                printer.Print(vars, "template class ClientReader<$outputTypeName$>;\n");
                printer.Print(vars,
                              "template Status ClientReader<$outputTypeName$>::Write(const $inputTypeName$ &);\n");
                readerClass.insert(outputName);
            }
        }
    }
}

void ZmqRpcGenerator::InstantiateTemplate(io::Printer &printer, const google::protobuf::FileDescriptor &file) const
{
    printer.PrintRaw("namespace datasystem {\n");
    if (!packageName.empty()) {
        printer.PrintRaw(UseNameSpace(packageName));
    }
    for (auto i = 0; i < file.service_count(); ++i) {
        auto *svc = file.service(i);
        InstantiateService(printer, *svc);
    }
    printer.PrintRaw("}\n");
}

void ZmqRpcGenerator::ImplementGenericStubApiDef(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc,
                                                 const std::string &indent)
{
    const std::string &svcName = svc.name();
    const std::string stub = svcName + "_Stub";
    ImplementGenericStubConstructor(printer, svc, indent);
    ImplementGenericStubOtherFuncDef(printer, stub);
    for (auto j = 0; j < svc.method_count(); ++j) {
        if (svc.method(j) == nullptr) {
            continue;
        }
        auto &method = *(svc.method(j));
        if (method.client_streaming() && method.server_streaming()) {
            ImplementGenericStubStreamingDef(printer, method, indent, stub);
        } else if (method.client_streaming() && !method.server_streaming()) {
            ImplementGenericStubClientStreamingDef(printer, method, indent, stub);
        } else if (!method.client_streaming() && method.server_streaming()) {
            ImplementGenericStubServerStreamingDef(printer, method, indent, stub);
        } else {
            ImplementGenericStubAsyncWriteDef(printer, method, indent, stub);
            ImplementGenericStubAsyncReadDef(printer, method, indent, stub);
            ImplementGenericStubNoStreamDef(printer, method, indent, stub);
            ImplementGenericStubNoStreamDef2(printer, method, indent, stub);
            ImplGenericStubNoStreamShortDef(printer, method, indent, stub);
        }
    }
}

void ZmqRpcGenerator::ImplementGenericStubOtherFuncDef(io::Printer &printer, const std::string &stub)
{
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    const std::string forgetRequest =
        "void $stub$::ForgetRequest(int64_t tagId) {\n"
        "    stub_->ForgetRequest(tagId);\n"
        "}\n";
    const std::string isPeerAlive =
        "bool $stub$::IsPeerAlive(uint32_t threshold) {\n"
        "    return stub_->IsPeerAlive(threshold);\n"
        "}\n";
    const std::string cacheSession =
        "void $stub$::CacheSession(bool cache) {\n"
        "    stub_->CacheSession(cache);\n"
        "}\n";
    const std::string getInitStatus =
        "Status $stub$::GetInitStatus() {\n"
        "    return stub_->GetInitStatus();\n"
        "}\n";
    printer.Print(vars, forgetRequest.c_str());
    printer.Print(vars, isPeerAlive.c_str());
    printer.Print(vars, cacheSession.c_str());
    printer.Print(vars, getInitStatus.c_str());
}

void ZmqRpcGenerator::ImplementGenericStubConstructor(io::Printer &printer,
                                                      const google::protobuf::ServiceDescriptor &svc,
                                                      const std::string &indent)
{
    (void)indent;
    const std::string &svcName = svc.name();
    const std::string stub = svcName + "_Stub";
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["svcName"] = svcName;
    std::string impl =
        "$stub$::$stub$(std::shared_ptr<::datasystem::RpcChannel> channel, int32_t timeoutMs) {\n"
        "    stub_ = std::make_unique<$svcName$_ZmqStub>(channel, timeoutMs);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementGenericStubNoStreamDef(io::Printer &printer,
                                                      const google::protobuf::MethodDescriptor &method,
                                                      const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["upperMethodName"] = StringToUpper(method.name());
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";
    vars["optSendPayload2"] = HasPayloadSendOption(method) ? ", payload" : "";
    vars["optRecvPayload2"] = HasPayloadRecvOption(method) ? ", recvBuffer" : "";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const ::datasystem::RpcOptions &opt, const $inputTypeName$& rq,\n"
        "                                          $outputTypeName$& reply$optSendPayload1$$optRecvPayload1$) {\n"
        "    ::datasystem::Status rc;\n";
    printer.Print(vars, impl.c_str());

    std::string impl3 =
        "    rc = stub_->$methodName$(opt, rq, reply$optSendPayload2$$optRecvPayload2$);\n"
        "    return rc;\n"
        "}\n";
    printer.Print(vars, impl3.c_str());
}

void ZmqRpcGenerator::ImplementGenericStubNoStreamDef2(io::Printer &printer,
                                                       const google::protobuf::MethodDescriptor &method,
                                                       const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", input.second" : "";
    vars["optSendPayload2"] = HasPayloadSendOption(method) ? ", true" : ", false";
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", output.second" : "";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const ::datasystem::RpcOptions &opt,\n"
        "\t\tstd::unique_ptr<datasystem::ClientUnaryWriterReader<$inputTypeName$,$outputTypeName$>> *out) {\n"
        "    ::datasystem::Status rc;\n";
    printer.Print(vars, impl.c_str());

    std::string impl2 =
        "    rc = stub_->$methodName$(opt, out);\n"
        "    return rc;\n"
        "}\n";
    printer.Print(vars, impl2.c_str());
}

void ZmqRpcGenerator::ImplGenericStubNoStreamShortDef(io::Printer &printer,
                                                      const google::protobuf::MethodDescriptor &method,
                                                      const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";
    vars["optSendPayload2"] = HasPayloadSendOption(method) ? ", payload" : "";
    vars["optRecvPayload2"] = HasPayloadRecvOption(method) ? ", recvBuffer" : "";
    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const $inputTypeName$& rq,\n"
        "                                          $outputTypeName$& reply$optSendPayload1$$optRecvPayload1$) {\n"
        "    return $methodName$(::datasystem::RpcOptions() , rq, reply$optSendPayload2$$optRecvPayload2$);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementGenericStubAsyncWriteDef(io::Printer &printer,
                                                        const google::protobuf::MethodDescriptor &method,
                                                        const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["optSendPayload2"] = HasPayloadSendOption(method) ? ", payload" : "";
    vars["upperMethodName"] = StringToUpper(method.name());
    std::string impl =
        "::datasystem::Status $stub$::$methodName$AsyncWrite(const ::datasystem::RpcOptions &opt,\n"
        "                                                    const $inputTypeName$& rq,\n"
        "                                                    int64_t &tagId$optSendPayload1$) {\n"
        "    ::datasystem::Status rc;\n";
    printer.Print(vars, impl.c_str());

    std::string impl2 =
        "    rc = stub_->$methodName$AsyncWrite(opt, rq, tagId$optSendPayload2$);\n"
        "    return rc;\n"
        "}\n";
    printer.Print(vars, impl2.c_str());
    std::string impl3 =
        "::datasystem::Status $stub$::$methodName$AsyncWrite(const $inputTypeName$& rq,\n"
        "                                                    int64_t &tagId$optSendPayload1$) {\n"
        "    return $methodName$AsyncWrite(::datasystem::RpcOptions(), rq, tagId$optSendPayload2$);\n"
        "}\n";
    printer.Print(vars, impl3.c_str());
}

void ZmqRpcGenerator::ImplementGenericStubAsyncReadDef(io::Printer &printer,
                                                       const google::protobuf::MethodDescriptor &method,
                                                       const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";
    vars["optRecvPayload2"] = HasPayloadRecvOption(method) ? ", recvBuffer" : "";
    vars["upperMethodName"] = StringToUpper(method.name());
    std::string impl =
        "::datasystem::Status $stub$::$methodName$AsyncRead(int64_t tagId, $outputTypeName$ &reply$optRecvPayload1$,\n"
        "                                                   ::datasystem::RpcRecvFlags flags) {\n"
        "    ::datasystem::Status rc;\n";

    printer.Print(vars, impl.c_str());
    std::string impl2 =
        "    rc = stub_->$methodName$AsyncRead(tagId, reply$optRecvPayload2$, flags);\n"
        "    return rc;\n"
        "}\n";
    printer.Print(vars, impl2.c_str());
}

void ZmqRpcGenerator::ImplementGenericStubStreamingDef(io::Printer &printer,
                                                       const google::protobuf::MethodDescriptor &method,
                                                       const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["payloadSendOption"] = HasPayloadSendOption(method) ? "true" : "false";
    vars["payloadRecvOption"] = HasPayloadRecvOption(method) ? "true" : "false";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(\n"
        "    const ::datasystem::RpcOptions &opt,\n"
        "    std::unique_ptr<datasystem::ClientWriterReader<$inputTypeName$, $outputTypeName$>> *out) {\n"
        "    ::datasystem::Status rc;\n"
        "    {\n"
        "        rc = stub_->$methodName$(opt, out);\n"
        "    }\n"
        "    return rc;\n"
        "}\n"
        "::datasystem::Status $stub$::$methodName$(\n"
        "    std::unique_ptr<datasystem::ClientWriterReader<$inputTypeName$, $outputTypeName$>> *out) {\n"
        "    return $methodName$(::datasystem::RpcOptions(), out);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementGenericStubClientStreamingDef(io::Printer &printer,
                                                             const google::protobuf::MethodDescriptor &method,
                                                             const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["payloadSendOption"] = HasPayloadSendOption(method) ? "true" : "false";
    vars["payloadRecvOption"] = HasPayloadRecvOption(method) ? "true" : "false";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const ::datasystem::RpcOptions &opt,\n"
        "                                          std::unique_ptr<datasystem::ClientWriter<$inputTypeName$>> *out) {\n"
        "    ::datasystem::Status rc;\n"
        "    {\n"
        "        rc = stub_->$methodName$(opt, out);\n"
        "    }\n"
        "    return rc;\n"
        "}\n"
        "::datasystem::Status $stub$::$methodName$(std::unique_ptr<datasystem::ClientWriter<$inputTypeName$>> *out) {\n"
        "    return $methodName$(::datasystem::RpcOptions(), out);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementGenericStubServerStreamingDef(io::Printer &printer,
                                                             const google::protobuf::MethodDescriptor &method,
                                                             const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["optSendPayload2"] = HasPayloadSendOption(method) ? ", payload" : "";
    vars["payloadSendOption"] = HasPayloadSendOption(method) ? "true" : "false";
    vars["payloadRecvOption"] = HasPayloadRecvOption(method) ? "true" : "false";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const ::datasystem::RpcOptions &opt,\n"
        "                                          std::unique_ptr<datasystem::ClientReader<$outputTypeName$>> *out,\n"
        "                                          const $inputTypeName$& rq$optSendPayload1$) {\n"
        "    ::datasystem::Status rc;\n";
    printer.Print(vars, impl.c_str());

    std::string impl2 =
        "    rc = stub_->$methodName$(opt, out, rq$optSendPayload2$);\n"
        "    return rc;\n"
        "}\n"
        "::datasystem::Status $stub$::$methodName$(std::unique_ptr<datasystem::ClientReader<$outputTypeName$>> *out,\n"
        "                                          const $inputTypeName$& rq$optSendPayload1$) {\n"
        "    return $methodName$(::datasystem::RpcOptions(), out, rq$optSendPayload2$);\n}\n";
    printer.Print(vars, impl2.c_str());
}

void ZmqRpcGenerator::ImplementStubApiDef(io::Printer &printer, const google::protobuf::ServiceDescriptor &svc,
                                          const std::string &indent)
{
    const std::string &svcName = svc.name();
    const std::string stub = svcName + "_ZmqStub";
    for (auto j = 0; j < svc.method_count(); ++j) {
        if (svc.method(j) == nullptr) {
            continue;
        }
        auto &method = *(svc.method(j));
        if (method.client_streaming() && method.server_streaming()) {
            ImplementStubStreamingDef(printer, method, j, indent, stub);
        } else if (method.client_streaming() && !method.server_streaming()) {
            ImplementStubClientStreamingDef(printer, method, j, indent, stub);
        } else if (!method.client_streaming() && method.server_streaming()) {
            ImplementStubServerStreamingDef(printer, method, j, indent, stub);
        } else {
            ImplementStubAsyncWriteDef(printer, method, j, indent, stub);
            ImplementStubAsyncReadDef(printer, method, j, indent, stub);
            ImplementStubNoStreamDef(printer, method, j, indent, stub);
            ImplementStubNoStreamDef2(printer, method, j, indent, stub);
            ImplStubNoStreamShortDef(printer, method, j, indent, stub);
        }
    }
}

void ZmqRpcGenerator::ImplementStubAsyncWriteDef(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                 int methodIndex, const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["inputTypeName"] = method.input_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["optSendPayload2"] = HasPayloadSendOption(method) ? ", payload" : ", {}";
    vars["optSendPayload3"] = HasPayloadSendOption(method) ? ", payload" : "";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$AsyncWrite(const ::datasystem::RpcOptions &opt,\n"
        "                                                    const $inputTypeName$& rq,\n"
        "                                                    int64_t &tagId$optSendPayload1$) {\n"
        "    auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "    ::datasystem::RpcOptions o(opt);\n"
        "    o.SetHWM(2);\n"
        "    return pimpl_->AsyncWrite(o, ServiceName(), methodObj, rq$optSendPayload2$, tagId);\n"
        "}\n"
        "::datasystem::Status $stub$::$methodName$AsyncWrite(const $inputTypeName$& rq,\n"
        "                                                    int64_t &tagId$optSendPayload1$) {\n"
        "    return $methodName$AsyncWrite(::datasystem::RpcOptions(), rq, tagId$optSendPayload3$);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementStubStreamingDef(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                int methodIndex, const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(\n"
        "    const ::datasystem::RpcOptions &opt,\n"
        "    std::unique_ptr<datasystem::ClientWriterReader<$inputTypeName$, $outputTypeName$>> *out) {\n"
        "    ::datasystem::Status rc;\n"
        "    auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "    std::shared_ptr<::datasystem::ZmqMsgQueRef> sock;\n"
        "    std::string workerId;\n"
        "    try {\n"
        "      rc = pimpl_->GetStreamPeer(ServiceName(), methodObj->MethodIndex(), opt, sock, workerId);\n"
        "    } catch (const std::bad_alloc &e) {\n"
        "      rc = ::datasystem::Status{ datasystem::StatusCode::K_OUT_OF_MEMORY, \"Out of memory\"};\n"
        "    }\n"
        "    if (rc.IsError()) { return rc; }\n"
        "    auto stream = std::make_unique<datasystem::ClientWriterReaderImpl<$inputTypeName$, $outputTypeName$>>(\n"
        "        std::move(sock), ServiceName(), methodObj->MethodIndex(), workerId,\n"
        "        methodObj->HasPayloadSendOption(), methodObj->HasPayloadRecvOption());\n"
        "    *out = std::make_unique<datasystem::ClientWriterReader<\n"
        "                            $inputTypeName$, $outputTypeName$>>(std::move(stream));\n"
        "    return rc;\n"
        "}\n"
        "::datasystem::Status $stub$::$methodName$(\n"
        "    std::unique_ptr<datasystem::ClientWriterReader<$inputTypeName$, $outputTypeName$>> *out) {\n"
        "    return $methodName$(::datasystem::RpcOptions(), out);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementStubClientStreamingDef(io::Printer &printer,
                                                      const google::protobuf::MethodDescriptor &method, int methodIndex,
                                                      const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const ::datasystem::RpcOptions &opt,\n"
        "                                          std::unique_ptr<datasystem::ClientWriter<$inputTypeName$>> *out) {\n"
        "    ::datasystem::Status rc;\n"
        "    auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "    std::shared_ptr<::datasystem::ZmqMsgQueRef> sock;\n"
        "    std::string workerId;\n"
        "    try {\n"
        "      rc = pimpl_->GetStreamPeer(ServiceName(), methodObj->MethodIndex(), opt, sock, workerId);\n"
        "    } catch (const std::bad_alloc &e) {\n"
        "      rc = ::datasystem::Status{ datasystem::StatusCode::K_OUT_OF_MEMORY, \"Out of memory\"};\n"
        "    }\n"
        "    if (rc.IsError()) { return rc; }\n"
        "    auto writer = std::make_unique<datasystem::ClientWriterImpl<$inputTypeName$>>(\n"
        "        std::move(sock), ServiceName(), methodObj->MethodIndex(), workerId,\n"
        "        methodObj->HasPayloadSendOption(), methodObj->HasPayloadRecvOption());\n"
        "    *out = std::make_unique<datasystem::ClientWriter<$inputTypeName$>>(std::move(writer));\n"
        "    return rc;\n"
        "}\n"
        "::datasystem::Status $stub$::$methodName$(std::unique_ptr<datasystem::ClientWriter<$inputTypeName$>> *out) {\n"
        "    return $methodName$(::datasystem::RpcOptions(), out);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementStubServerStreamingDef(io::Printer &printer,
                                                      const google::protobuf::MethodDescriptor &method, int methodIndex,
                                                      const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["optSendPayload2"] = HasPayloadSendOption(method) ? ", payload" : "";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const ::datasystem::RpcOptions &opt,\n"
        "                                          std::unique_ptr<datasystem::ClientReader<$outputTypeName$>> *out,\n"
        "                                          const $inputTypeName$& rq$optSendPayload1$) {\n"
        "    ::datasystem::Status rc;\n"
        "    auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "    std::shared_ptr<::datasystem::ZmqMsgQueRef> sock;\n"
        "    std::string workerId;\n"
        "    try {\n"
        "      rc = pimpl_->GetStreamPeer(ServiceName(), methodObj->MethodIndex(), opt, sock, workerId);\n"
        "    } catch (const std::bad_alloc &e) {\n"
        "      rc = ::datasystem::Status{ datasystem::StatusCode::K_OUT_OF_MEMORY, \"Out of memory\"};\n"
        "    }\n"
        "    if (rc.IsError()) { return rc; }\n"
        "    auto reader = std::make_unique<::datasystem::ClientReaderImpl<$outputTypeName$>>(\n"
        "        std::move(sock), ServiceName(), methodObj->MethodIndex(), workerId,\n"
        "        methodObj->HasPayloadSendOption(), methodObj->HasPayloadRecvOption());\n"
        "    rc = reader->WritePb(rq);\n"
        "    if (rc.IsError()) { return rc; }\n";
    if (HasPayloadSendOption(method)) {
        impl +=
            "    rc = reader->SendPayload(payload);\n"
            "    if (rc.IsError()) { return rc; }\n";
    }
    impl +=
        "    *out = std::make_unique<datasystem::ClientReader<$outputTypeName$>>(std::move(reader));\n"
        "    return rc;\n"
        "}\n"
        "::datasystem::Status $stub$::$methodName$(std::unique_ptr<datasystem::ClientReader<$outputTypeName$>> *out,\n"
        "                                          const $inputTypeName$& rq$optSendPayload1$) {\n"
        "    return $methodName$(::datasystem::RpcOptions(), out, rq$optSendPayload2$);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementStubAsyncReadDef(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                int methodIndex, const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["outputTypeName"] = method.output_type()->name();
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$AsyncRead(int64_t tagId, $outputTypeName$ &reply$optRecvPayload1$,\n"
        "                                                   ::datasystem::RpcRecvFlags flags) {\n"
        "    auto &methodObj = methodMap_.find($methodIndex$)->second;\n";
    if (!HasPayloadRecvOption(method)) {
        impl += "    std::vector<::datasystem::RpcMessage> recvBuffer;\n";
    }
    impl +=
        "    return pimpl_->AsyncRead(tagId, ServiceName(), methodObj, reply, recvBuffer, flags);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementStubNoStreamDef(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                               int methodIndex, const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["upperMethodName"] = StringToUpper(method.name());
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";
    std::string impl;
    ImplementStubNoStreamDefHelper(impl, method);
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplementStubNoStreamDefHelper(std::string &impl,
                                                     const google::protobuf::MethodDescriptor &method)
{
    impl =
        "::datasystem::Status $stub$::$methodName$(const ::datasystem::RpcOptions &opt, const $inputTypeName$& rq,\n"
        "                                          $outputTypeName$& reply$optSendPayload1$$optRecvPayload1$) {\n";
    bool criticalFunc = CRITICAL_FUNCTIONS.find(method.name()) != CRITICAL_FUNCTIONS.end();
    if (criticalFunc) {
        impl += "    PerfPoint point1(PerfKey::ZMQ_ADAPTOR_LAYER_$upperMethodName$);\n";
    }
    impl +=
        "    ::datasystem::Status rc;\n"
        "    auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "    std::shared_ptr<::datasystem::ZmqMsgQueRef> sock;\n"
        "    ::datasystem::RpcOptions o(opt);\n"
        "    o.SetHWM(2);\n"
        "    rc = pimpl_->CreateMsgQ(sock, serviceName_, o);\n"
        "    if (rc.IsError()) { return rc; }\n";

    if (criticalFunc) {
        impl +=
            "    point1.Record();\n"
            "    PerfPoint point2(PerfKey::ZMQ_$upperMethodName$_RPC);\n";
    }
    impl +=
        "    auto clientApi =\n"
        "        std::make_unique<datasystem::ClientUnaryWriterReaderImpl<$inputTypeName$, $outputTypeName$>>(\n"
        "            std::move(sock), ServiceName(), methodObj->MethodIndex(),\n"
        "            methodObj->HasPayloadSendOption(), methodObj->HasPayloadRecvOption());\n"
        "    rc = clientApi->Write(rq);\n"
        "    if (rc.IsError()) { return rc; }\n";

    if (HasPayloadSendOption(method)) {
        impl +=
            "    rc = clientApi->SendPayload(payload);\n"
            "    if (rc.IsError()) { return rc; }\n";
    }
    impl +=
        "    rc = clientApi->Read(reply);\n"
        "    if (rc.IsError()) { return rc; }\n";

    if (HasPayloadRecvOption(method)) {
        impl +=
            "    rc = clientApi->ReceivePayload(recvBuffer);\n"
            "    if (rc.IsError()) { return rc; }\n";
    }
    impl +=
        "    return rc;\n"
        "}\n";
}

void ZmqRpcGenerator::ImplementStubNoStreamDef2(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                                int methodIndex, const std::string &indent, const std::string &stub)
{
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const ::datasystem::RpcOptions &opt,\n"
        "\t\tstd::unique_ptr<datasystem::ClientUnaryWriterReader<$inputTypeName$, $outputTypeName$>> *out) {\n"
        "    ::datasystem::Status rc;\n"
        "    auto &methodObj = methodMap_.find($methodIndex$)->second;\n"
        "    std::shared_ptr<::datasystem::ZmqMsgQueRef> sock;\n"
        "    ::datasystem::RpcOptions o(opt);\n"
        "    o.SetHWM(2);\n"
        "    rc = pimpl_->CreateMsgQ(sock, serviceName_, o);\n"
        "    if (rc.IsError()) { return rc; }\n"
        "    auto clientApi =\n"
        "        std::make_unique<datasystem::ClientUnaryWriterReaderImpl<$inputTypeName$, $outputTypeName$>>(\n"
        "            std::move(sock), ServiceName(), methodObj->MethodIndex(),\n"
        "            methodObj->HasPayloadSendOption(), methodObj->HasPayloadRecvOption());\n"
        "    *out = std::make_unique<datasystem::ClientUnaryWriterReader<$inputTypeName$, "
        "$outputTypeName$>>(std::move(clientApi));\n"
        "    return rc;\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

void ZmqRpcGenerator::ImplStubNoStreamShortDef(io::Printer &printer, const google::protobuf::MethodDescriptor &method,
                                               int methodIndex, const std::string &indent, const std::string &stub)
{
    (void)methodIndex;
    (void)indent;
    std::map<std::string, std::string> vars;
    vars["stub"] = stub;
    vars["methodName"] = method.name();
    vars["methodIndex"] = std::to_string(methodIndex);
    vars["inputTypeName"] = method.input_type()->name();
    vars["outputTypeName"] = method.output_type()->name();
    vars["optSendPayload1"] = HasPayloadSendOption(method) ? ", const std::vector<::datasystem::MemView> &payload" : "";
    vars["optRecvPayload1"] = HasPayloadRecvOption(method) ? ", std::vector<::datasystem::RpcMessage> &recvBuffer" : "";
    vars["optSendPayload2"] = HasPayloadSendOption(method) ? ", payload" : "";
    vars["optRecvPayload2"] = HasPayloadRecvOption(method) ? ", recvBuffer" : "";

    std::string impl =
        "::datasystem::Status $stub$::$methodName$(const $inputTypeName$& rq,\n"
        "                                          $outputTypeName$& reply$optSendPayload1$$optRecvPayload1$) {\n"
        "    return $methodName$(::datasystem::RpcOptions() "
        ", rq, reply$optSendPayload2$$optRecvPayload2$);\n"
        "}\n";
    printer.Print(vars, impl.c_str());
}

}  // namespace datasystem
